/*
 * This is written by Lulu Han.
 * E-mail: locomotive_crypto@163.com
 */

#include <iostream>
#include <exception>
#include "sec_sed.h"
#include "timer.h"
#include "net_client.h"
#include "net_server.h"

void test_sed_for_correctness(long bitlens);

void test_sed_for_performance(long bits, long cnt);

void test_sed_by_network_SA(long bits, long cnt, size_t port, std::string ipaddr);

void test_sed_by_network_SB(size_t port, long vlens);

int main(int argc, char* argv[]){
	
	std::cout << "----------------------------------------------------------" << std::endl;
	std::cout << "-----------------------<Usage>----------------------------" << std::endl;
	std::cout << "-----1) Run for correctness: ./program--------------------" << std::endl;
	std::cout << "-----2) Run for performance: ./program cnt----------------" << std::endl;
	std::cout << "-----3) Run as SB: ./program port ------------------------" << std::endl;
	std::cout << "-----4) Run as SA: ./program port ipaddr------------------" << std::endl;
	std::cout << "----------------------------------------------------------" << std::endl << std::endl;
	
	if(argc == 1){
		test_sed_for_correctness(200);
	}
	else if(argc == 2){
		for(size_t i = 0; i < 5; ++i)
			test_sed_for_performance(512, std::atoi(argv[1]));
		std::cout << std::endl;
		
		for(size_t i = 0; i < 5; ++i)
			test_sed_for_performance(1024, std::atoi(argv[1]));
		std::cout << std::endl;
		
		for(size_t i = 0; i < 5; ++i)
			test_sed_for_performance(1536, std::atoi(argv[1]));
		std::cout << std::endl;
	}
	else if(argc == 3){
		if(std::string(argv[1]) == std::string("sb") || std::string(argv[1]) == std::string("SB")){
			test_sed_by_network_SB(std::atoi(argv[2]), 10);
		}
			
	}
	else if(argc == 5){
		if(std::string(argv[1]) == std::string("sa") || std::string(argv[1]) == std::string("SA")){
			//test_sed_by_network_SA(512, std::atoi(argv[4]), std::atoi(argv[2]), std::string(argv[3]));
			
			//test_sed_by_network_SA(1024, std::atoi(argv[4]), std::atoi(argv[2]), std::string(argv[3]));
			
			test_sed_by_network_SA(1536, std::atoi(argv[4]), std::atoi(argv[2]), std::string(argv[3]));
		}
	}
	
	return 0;
}


void test_sed_for_correctness(long bitlens){
	PrivateKey prikey;
	PublicKey pubkey;
	
	Paillier::GenKeyPair(prikey, pubkey, bitlens);
	
	Decryptor dec(prikey);
	Encryptor enc(pubkey);
	
	NTL::ZZ a0, a1;
	NTL::ZZ b0, b1;
	
	a0 = 123; a1 = 567;
	b0 = 901; b1 = 345;
	
	std::cout << "a0 = " << a0 << std::endl;
	std::cout << "a1 = " << a1 << std::endl;
	std::cout << "b0 = " << b0 << std::endl;
	std::cout << "b1 = " << b1 << std::endl;
	
	NTL::ZZ ca0, ca1;
	NTL::ZZ cb0, cb1;
	
	enc.Encrypt(a0, ca0);
	enc.Encrypt(a1, ca1);
	enc.Encrypt(b0, cb0);
	enc.Encrypt(b1, cb1);
	
	std::vector<NTL::ZZ> cva, cvb;
	
	cva.push_back(ca0);
	cva.push_back(ca1);
	
	cvb.push_back(cb0);
	cvb.push_back(cb1);
	
	SecSed::CloudA CA(prikey);
	SecSed::CloudB CB(pubkey);
	
	std::vector<NTL::ZZ> ncva;
	std::vector<NTL::ZZ> ncvb;
	CB.gen_two_ciphers(cva, cvb, ncva, ncvb);
	
	std::vector<NTL::ZZ> cvmul;
	CA.dec_reenc(ncva, ncvb, cvmul);
	
	NTL::ZZ csed;
	CB.gen_mul(cvmul, csed);
	std::cout << "csed = " << csed << std::endl;
	
	NTL::ZZ sed;
	dec.Decrypt(csed, sed);
	std::cout << "sed = " << sed << std::endl;
	
	if(sed != (((a0-b0)*(a0-b0) + (a1-b1)*(a1-b1)) % enc.GetN())){
		throw std::runtime_error("Failed to test SecSed ...\n");
	}
}

NTL::ZZ dist(const std::vector<NTL::ZZ>& va, const std::vector<NTL::ZZ>& vb){
	NTL::ZZ da, db;
	
	da = (va[0] - vb[0]) * (va[0] - vb[0]);
	db = (va[1] - vb[1]) * (va[1] - vb[1]);
	
	return (da + db);
}

void test_sed_for_performance(long bits, long cnt){
	std::cout << "The key length for this test is: " << 2*bits << std::endl;
	 
	std::vector<Paillier::Encryptor> vencryptor;
	std::vector<Paillier::Decryptor> vdecryptor;
	
	std::vector<SecSed::CloudA> vCA;
	std::vector<SecSed::CloudB> vCB;
	
	for(long i = 0; i < cnt; ++i){
		Paillier::PrivateKey prikey;
		Paillier::PublicKey pubkey;
	
		// Generate random key pair
		Paillier::GenKeyPair(prikey, pubkey, bits);
	
		// Define Encryptor and Decryptor objects
		Paillier::Encryptor encryptor(pubkey);
		Paillier::Decryptor decryptor(prikey);
		
		vencryptor.push_back(encryptor);
		vdecryptor.push_back(decryptor);
		
		SecSed::CloudA CA(prikey);
		SecSed::CloudB CB(pubkey);
		
		vCA.push_back(CA);
		vCB.push_back(CB);
	}
	
	std::vector<std::vector<NTL::ZZ>> vcva, vcvb;
	std::vector<std::vector<NTL::ZZ>> vva, vvb;
	NTL::ZZ r;
	NTL::ZZ a0, a1;
	NTL::ZZ b0, b1;
	NTL::ZZ ca0, ca1;
	NTL::ZZ cb0, cb1;
	for(long i = 0; i < cnt; ++i){
		std::vector<NTL::ZZ> va, vb;
		
		a0 = NTL::RandomLen_ZZ(30);
		a1 = NTL::RandomLen_ZZ(30);
		
		va.push_back(a0);
		va.push_back(a1);
		vva.push_back(va);
		
		b0 = NTL::RandomLen_ZZ(30);
		b1 = NTL::RandomLen_ZZ(30);
		
		vb.push_back(b0);
		vb.push_back(b1);
		vvb.push_back(vb);
		
		std::vector<NTL::ZZ> cva, cvb;
		
		vencryptor[i].Encrypt(a0, ca0);
		vencryptor[i].Encrypt(a1, ca1);
		
		cva.push_back(ca0);
		cva.push_back(ca1);
		
		vencryptor[i].Encrypt(b0, cb0);
		vencryptor[i].Encrypt(b1, cb1);
		
		cvb.push_back(cb0);
		cvb.push_back(cb1);
		
		vcva.push_back(cva);
		vcvb.push_back(cvb);
	}
	
	time_point<high_resolution_clock> start;	
	start = now_time();
	for(long i = 0; i < cnt; ++i){
		std::vector<NTL::ZZ> ncva;
		std::vector<NTL::ZZ> ncvb;
		vCB[i].gen_two_ciphers(vcva[i], vcvb[i], ncva, ncvb);
	
		std::vector<NTL::ZZ> cvmul;
		vCA[i].dec_reenc(ncva, ncvb, cvmul);
	
		NTL::ZZ csed;
		vCB[i].gen_mul(cvmul, csed);
		
		/*
		NTL::ZZ psed;
		vdecryptor[i].Decrypt(csed, psed);
		
		NTL::ZZ fsed = dist(vva[i], vvb[i]) % vencryptor[i].GetN();
		std::cout << "psed = " << psed << std::endl;
		std::cout << "fsed = " << fsed << std::endl;
		
		if(psed != fsed){
			throw std::runtime_error("Failed to test SecSed ...\n");
		}*/
	}
	double tm = time_from(start);
	std::cout << "The time for "<< cnt << " SecSed is : " << tm << std::endl;
}

void test_sed_by_network_SA(long bits, long cnt, size_t port, std::string ipaddr){
	std::cout << "The key length for this test is: " << 2*bits << std::endl;
	
	NetClient client(ipaddr, port);
	client.connect();
	
	client.send(cnt);
	client.send(bits);
	
	std::vector<Paillier::Decryptor> vdecryptor;
	
	std::vector<SecSed::CloudA> vCA;
	
	for(long i = 0; i < cnt; ++i){
		Paillier::PrivateKey prikey;
		Paillier::PublicKey pubkey;
	
		// Generate random key pair
		Paillier::GenKeyPair(prikey, pubkey, bits);
		
		client.send(pubkey.GetN());
		
		// Define Encryptor and Decryptor objects
		Paillier::Decryptor decryptor(prikey);
		
		vdecryptor.push_back(decryptor);
		
		SecSed::CloudA CA(prikey);
		
		vCA.push_back(CA);
	}
	
	std::cout << "Please count the network traffics from here..." << std::endl;
	getchar();
	for(long i = 0; i < cnt; ++i){
		std::vector<NTL::ZZ> ncva;
		std::vector<NTL::ZZ> ncvb;
		client.receive(ncva);
		client.receive(ncvb);
		
		std::vector<NTL::ZZ> cvmul;
		vCA[i].dec_reenc(ncva, ncvb, cvmul);
		client.send(cvmul);
		
		/*
		NTL::ZZ csed;
		client.receive(csed);	
		
		NTL::ZZ psed;
		vdecryptor[i].Decrypt(csed, psed);
		std::cout << "psed = " << psed << std::endl;
		*/
	}
	
	client.close();
}

void test_sed_by_network_SB(size_t port, long vlens){
	NetServer server(port); 
	server.accept();
	
	long cnt;
	server.receive(cnt);
	
	long bits;
	server.receive(bits);
	
	std::cout << "The key length for this test is: " << 2*bits << std::endl;
	
	std::vector<Paillier::Encryptor> vencryptor;
	std::vector<SecSed::CloudB> vCB;
	
	NTL::ZZ N;
	for(long i = 0; i < cnt; ++i){
		server.receive(N);
		
		Paillier::PublicKey pubkey(N);
		SecSed::CloudB CB(pubkey);
		
		Paillier::Encryptor enc(pubkey);
		vencryptor.push_back(enc);
		
		vCB.push_back(CB);
	}
	
	std::vector<std::vector<NTL::ZZ>> vcva, vcvb;
	std::vector<std::vector<NTL::ZZ>> vva, vvb;
	NTL::ZZ r;
	NTL::ZZ a0, a1;
	NTL::ZZ b0, b1;
	NTL::ZZ ca0, ca1;
	NTL::ZZ cb0, cb1;
	for(long i = 0; i < cnt; ++i){
		std::vector<NTL::ZZ> va, vb;
		
		a0 = NTL::RandomLen_ZZ(30);
		a1 = NTL::RandomLen_ZZ(30);
		
		va.push_back(a0);
		va.push_back(a1);
		vva.push_back(va);
		
		b0 = NTL::RandomLen_ZZ(30);
		b1 = NTL::RandomLen_ZZ(30);
		
		vb.push_back(b0);
		vb.push_back(b1);
		vvb.push_back(vb);
		
		std::vector<NTL::ZZ> cva, cvb;
		
		vencryptor[i].Encrypt(a0, ca0);
		vencryptor[i].Encrypt(a1, ca1);
		
		cva.push_back(ca0);
		cva.push_back(ca1);
		
		vencryptor[i].Encrypt(b0, cb0);
		vencryptor[i].Encrypt(b1, cb1);
		
		cvb.push_back(cb0);
		cvb.push_back(cb1);
		
		vcva.push_back(cva);
		vcvb.push_back(cvb);
	}
	
	std::cout << "Please count the network traffics from here..." << std::endl;
	getchar();
	for(long i = 0; i < cnt; ++i){
		std::vector<NTL::ZZ> ncva;
		std::vector<NTL::ZZ> ncvb;
		vCB[i].gen_two_ciphers(vcva[i], vcvb[i], ncva, ncvb);
		server.send(ncva);
		server.send(ncvb);
		
		std::vector<NTL::ZZ> cvmul;
		server.receive(cvmul);
		
		NTL::ZZ csed;
		vCB[i].gen_mul(cvmul, csed);
		//server.send(csed);
		
		//NTL::ZZ fsed = dist(vva[i], vvb[i]) % vencryptor[i].GetN();
		//std::cout << "fsed = " << fsed << std::endl;
	}

	server.close();
}
